#include "kernel_operator.h"
#include <type_traits>
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2;

template<typename T> class KernelTril {
public:
    __aicore__ inline KernelTril() {}
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, int64_t M, int64_t N, int64_t Z, int64_t diagonal) {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");

        this->M = M; //高
        this->N = N; //长
        this->Z = Z; //循环次数
        this->diagonal = diagonal;
        this->row32 = (N+31)/32*32;

        xGm.SetGlobalBuffer(reinterpret_cast<__gm__ DTYPE_X *>(x), M * N * Z);
        yGm.SetGlobalBuffer(reinterpret_cast<__gm__ DTYPE_Y *>(y), M * N * Z);

        pipe.InitBuffer(inQueueX, BUFFER_NUM, (N+31)/32*32 * sizeof(DTYPE_X));
        pipe.InitBuffer(outQueueY, BUFFER_NUM, (N+31)/32*32 * sizeof(DTYPE_Y));
        pipe.InitBuffer(QueueZero, (N+31)/32*32 * sizeof(DTYPE_X));
    }
    __aicore__ inline void Process() {

        for(int32_t i = 0; i < Z; i ++)
        {
            for(int32_t j = 0; j < M; j++)
            {
                int32_t num0 = j + this->diagonal + 1;
                if(num0 < 0) num0 = 0;
                if(num0 > N) num0 = N;
                CopyIn(i, j, num0);
                Compute(i, j, num0);
                CopyOut(i, j);
            }
        }
    }
private:
     __aicore__ inline void CopyIn(int32_t i, int32_t j, int32_t num0)
    {
        LocalTensor<DTYPE_X> xLocal = inQueueX.AllocTensor<DTYPE_X>();
        if(num0 > 0)
            DataCopy(xLocal, xGm[i*M*N + j*N], (num0+31)/32*32);
        inQueueX.EnQue(xLocal);
    }
    __aicore__ inline void Compute(int32_t i, int32_t j, int32_t num0)
    {
        LocalTensor<DTYPE_X> xLocal = inQueueX.DeQue<DTYPE_X>();
        LocalTensor<DTYPE_Y> yLocal = outQueueY.AllocTensor<DTYPE_Y>();

        Duplicate(yLocal, static_cast<DTYPE_X>(0), N);

        if(num0 > 0)
            Add(yLocal, yLocal, xLocal, num0);

        outQueueY.EnQue<DTYPE_Y>(yLocal);
        inQueueX.FreeTensor(xLocal);
    }
    __aicore__ inline void CopyOut(int32_t i, int32_t j)
    {
        LocalTensor<DTYPE_Y> yLocal = outQueueY.DeQue<DTYPE_Y>();
        DataCopy(yGm[i*M*N + j*N], yLocal, row32);
        outQueueY.FreeTensor(yLocal);
    }
private:
    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueY;
    TBuf<QuePosition::VECCALC> QueueZero;

    GlobalTensor<DTYPE_X> xGm;
    GlobalTensor<DTYPE_Y> yGm;

    int64_t M;
    int64_t N;
    int64_t Z;
    int64_t diagonal;
    int64_t row32;
};

extern "C" __global__ __aicore__ void tril(GM_ADDR x, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    // TODO: user kernel impl
    KernelTril<DTYPE_X> op;
    op.Init(x, y, tiling_data.M, tiling_data.N, tiling_data.Z, tiling_data.diagonal);
    op.Process();
}